package cn.fzkj.ws;

import cn.fzkj.ws.message.WsMessage;
import com.alibaba.fastjson2.JSON;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 * @author mr
 * @description websocket消息处理器
 * @since 2024/2/20
 */
@Slf4j
public class SocketHandler extends TextWebSocketHandler {

    /**
     * 会话 <userId, WebSocketSession>
     */
    protected final Map<String, List<WebSocketSession>> sessions = new ConcurrentHashMap<>();

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        String msg = message.getPayload();
        if ("ping".equalsIgnoreCase(msg)) {
            session.sendMessage(new TextMessage("pong"));
            return;
        }
        // 处理其他消息
        log.info("websocket 接收消息{}, 消息：{}", session.getId(), msg);
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("websocket 建立连接 {}", session.getId());
        String userId = getUserId(session);
        // 一个用户可能在多个浏览器登录，有多个session,需要保存一个用户多session
        if (!sessions.containsKey(userId)) {
            sessions.put(userId, new ArrayList<>());
        }
        sessions.get(userId).add(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("websocket 连接 {} 关闭 原因：{}", session.getId(), status.getReason());
        if (session.isOpen()) {
            session.close();
        }
        String userId = getUserId(session);
        sessions.remove(userId);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable throwable) throws Exception {
        log.error("websocket 异常 {} {}", session, throwable.getMessage());
        String userId = getUserId(session);
        sessions.remove(userId);
    }

    /**
     * 发送消息
     *
     * @param msg msg
     */
    public void sendMessage(WsMessage msg) {
        if (msg == null) {
            return;
        }
        TextMessage message = new TextMessage(JSON.toJSONString(msg));
        String to = msg.getTo();
        List<WebSocketSession> userSessions = sessions.get(to);
        if (CollectionUtils.isEmpty(userSessions)) {
            log.warn("用户{}不在线", to);
            return;
        }
        userSessions.forEach(webSocketSession -> {
            if (webSocketSession == null || !webSocketSession.isOpen()) {
                log.warn("用户{}不在线", to);
                return;
            }
            try {
                webSocketSession.sendMessage(message);
            } catch (IOException e) {
                log.error("发送消息给{}失败，异常{}", to, e.getMessage());
            }
        });

    }

    /**
     * 发送消息 广泛播
     *
     * @param tos TOS
     * @param msg msg
     */
    public void sendMessageBroadCast(List<String> tos, WsMessage msg) {
        if (CollectionUtils.isEmpty(tos)) {
            return;
        }
        for (String to : tos) {
            msg.setTo(to);
            sendMessage(msg);
        }
    }

    private String getUserId(WebSocketSession session) {
        Map<String, Object> attributes = session.getAttributes();
        Object userId = attributes.get("userId");
        if (userId == null) {
            throw new RuntimeException("无效的websocket连接");
        }
        return (String) userId;
    }

}
